from Crypto.Cipher import DES, DES3
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import base64

class EncryptionError(Exception):
    pass

class DecryptionError(Exception):
    pass

class DESCipher:
    """DES/3DES加密实现，支持ECB/CBC/CTR/OFB/CFB模式。"""
    # 添加加密模式常量
    MODE_ECB = DES.MODE_ECB
    MODE_CBC = DES.MODE_CBC
    MODE_CTR = DES.MODE_CTR
    MODE_OFB = DES.MODE_OFB
    MODE_CFB = DES.MODE_CFB
    
    # 修改CTR模式的最大块大小（以字节为单位）
    CTR_MAX_BLOCK_SIZE = 8 * 1024  # 8KB
    
    def __init__(self, key, triple_des=False, mode=DES.MODE_CBC, iv=None, nonce=None):
        """
        使用密钥和模式初始化DES/3DES加密器。
        
        参数:
            key (bytes): 加密密钥（DES为8字节，3DES为24字节）
            triple_des (bool): 是否使用三重DES
            mode (int): 加密模式（DES.MODE_CBC、DES.MODE_CTR、DES.MODE_OFB、DES.MODE_CFB或DES.MODE_ECB）
            iv (bytes): 初始化向量（对于CBC、OFB、CFB模式必需）
            nonce (bytes): 随机数（对于CTR模式必需）
        """
        self.triple_des = triple_des
        self.mode = mode
        self.key = self._process_key(key)
        self.cipher_class = DES3 if triple_des else DES
        
        if mode == DES.MODE_CTR:
            if nonce:
                if isinstance(nonce, str):
                    nonce = base64.b64decode(nonce)
                if len(nonce) > 7:  # DES CTR模式nonce最大长度为7字节
                    raise ValueError("CTR模式nonce长度不能超过7字节")
            else:
                nonce = get_random_bytes(7)  # 使用7字节nonce
            self.nonce = nonce
            self.iv = None
        elif mode != DES.MODE_ECB:
            if iv:
                if isinstance(iv, str):
                    iv = base64.b64decode(iv)
                if len(iv) != 8:  # DES块大小为8字节
                    raise ValueError("IV长度必须为8字节")
            else:
                iv = get_random_bytes(8)
            self.iv = iv
            self.nonce = None
        else:
            self.iv = None
            self.nonce = None
    
    def _process_key(self, key):
        """确保密钥满足长度要求。"""
        if not key:
            raise ValueError("密钥不能为空")
            
        if isinstance(key, str):
            key = key.encode('utf-8')
        
        required_length = 24 if self.triple_des else 8
        if len(key) != required_length:
            if self.triple_des:
                # 对于3DES，如果密钥长度不足24字节，使用密钥派生
                if len(key) < 8:
                    key = key.ljust(8, b'\0')
                # 使用密钥派生生成三个不同的8字节密钥部分
                key1 = key[:8]
                key2 = key1[::-1]  # 反转key1作为key2
                key3 = key1[::2] + key1[1::2]  # 交错key1的字节作为key3
                key = key1 + key2 + key3
            else:
                key = key.ljust(required_length, b'\0')[:required_length]
        
        # 验证3DES密钥的三个部分是否相同
        if self.triple_des:
            key1, key2, key3 = key[:8], key[8:16], key[16:24]
            if key1 == key2 or key2 == key3 or key1 == key3:
                raise ValueError("3DES密钥的三个部分不能相同")
                
        return key
    
    def encrypt(self, data):
        """
        使用DES/3DES加密数据。
        
        参数:
            data: 要加密的字符串或字节
            
        返回:
            tuple: (Base64编码的加密数据, Base64编码的IV/nonce)
            
        异常:
            EncryptionError: 如果加密失败
        """
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            
            if self.mode == self.cipher_class.MODE_CTR:
                # 对于CTR模式，如果数据较大，进行分块处理
                if len(data) > self.CTR_MAX_BLOCK_SIZE:
                    encrypted_blocks = []
                    total_blocks = (len(data) + self.CTR_MAX_BLOCK_SIZE - 1) // self.CTR_MAX_BLOCK_SIZE
                    
                    # 生成基础nonce
                    base_nonce = get_random_bytes(6)
                    
                    for i in range(0, len(data), self.CTR_MAX_BLOCK_SIZE):
                        block = data[i:i + self.CTR_MAX_BLOCK_SIZE]
                        # 为每个块生成新的nonce，使用块索引作为nonce的一部分
                        block_index = i // self.CTR_MAX_BLOCK_SIZE
                        # 将块索引转换为字节并添加到nonce中
                        block_nonce = base_nonce + bytes([block_index & 0xFF])
                        
                        # 创建新的加密器实例
                        cipher = self.cipher_class.new(self.key, self.mode, nonce=block_nonce)
                        encrypted_block = cipher.encrypt(block)
                        # 将nonce和加密数据一起存储
                        encrypted_blocks.append(base64.b64encode(block_nonce + encrypted_block).decode('utf-8'))
                    
                    # 使用特殊分隔符连接所有块
                    result = '||'.join(encrypted_blocks)
                    return result, None
                else:
                    # 小数据块直接处理
                    cipher = self.cipher_class.new(self.key, self.mode, nonce=self.nonce)
                    encrypted_data = cipher.encrypt(data)
                    result = base64.b64encode(encrypted_data).decode('utf-8')
                    iv_str = base64.b64encode(self.nonce).decode('utf-8')
                    return result, iv_str
            elif self.mode == self.cipher_class.MODE_ECB:
                cipher = self.cipher_class.new(self.key, self.mode)
                iv = None
            else:
                cipher = self.cipher_class.new(self.key, self.mode, self.iv)
                iv = self.iv
            
            if self.mode == self.cipher_class.MODE_CTR:
                encrypted_data = cipher.encrypt(data)  # CTR模式不需要填充
            else:
                padded_data = pad(data, 8)  # DES块大小为8字节
                encrypted_data = cipher.encrypt(padded_data)
            
            result = base64.b64encode(encrypted_data).decode('utf-8')
            iv_str = base64.b64encode(iv).decode('utf-8') if iv else None
            
            return result, iv_str
            
        except Exception as e:
            raise EncryptionError(f"{'3DES' if self.triple_des else 'DES'}加密失败: {str(e)}")

    def decrypt(self, encrypted_data, iv=None):
        """
        解密DES/3DES加密的数据。
        
        参数:
            encrypted_data: Base64编码的加密数据
            iv: Base64编码的IV/nonce（对于CBC、CTR、OFB、CFB模式必需）
            
        返回:
            解密后的字符串
            
        异常:
            DecryptionError: 如果解密失败
        """
        try:
            if self.mode == self.cipher_class.MODE_CTR:
                # 检查是否是分块加密的数据
                if '||' in encrypted_data:
                    decrypted_blocks = []
                    for block in encrypted_data.split('||'):
                        block_bytes = base64.b64decode(block)
                        block_nonce = block_bytes[:7]
                        encrypted_block = block_bytes[7:]
                        cipher = self.cipher_class.new(self.key, self.mode, nonce=block_nonce)
                        decrypted_block = cipher.decrypt(encrypted_block)
                        decrypted_blocks.append(decrypted_block)
                    return b''.join(decrypted_blocks).decode('utf-8')
                else:
                    encrypted_bytes = base64.b64decode(encrypted_data)
                    if not iv:
                        raise ValueError("CTR模式需要提供nonce")
                    nonce_bytes = base64.b64decode(iv)
                    if len(nonce_bytes) > 7:
                        raise ValueError("CTR模式nonce长度不能超过7字节")
                    cipher = self.cipher_class.new(self.key, self.mode, nonce=nonce_bytes)
                    decrypted_data = cipher.decrypt(encrypted_bytes)
                    return decrypted_data.decode('utf-8')
            else:
                encrypted_bytes = base64.b64decode(encrypted_data)
                
                if self.mode == self.cipher_class.MODE_ECB:
                    cipher = self.cipher_class.new(self.key, self.mode)
                else:
                    if not iv:
                        raise ValueError("非ECB模式需要提供IV")
                    iv_bytes = base64.b64decode(iv)
                    cipher = self.cipher_class.new(self.key, self.mode, iv_bytes)
                
                if self.mode == self.cipher_class.MODE_CTR:
                    decrypted_data = cipher.decrypt(encrypted_bytes)
                else:
                    decrypted_data = unpad(cipher.decrypt(encrypted_bytes), 8)
                
                return decrypted_data.decode('utf-8')
            
        except Exception as e:
            raise DecryptionError(f"{'3DES' if self.triple_des else 'DES'}解密失败: {str(e)}")